import heapq
import numpy as np
import cv2
from pathlib import Path
from opencv import draw_contours_to_mask, cv_show

__all__ = ['get_largest_cnt_from_mask', 'segment_largest_roi_from_image', 'find_bounding_rect', 'resize_uniform_same_size', 'remove_small_objects']


def erode(image, ks=5, iters=2):
    kernel = np.ones((ks, ks), np.uint8)
    erode_image = cv2.erode(image, kernel, iters)
    return erode_image


def dilate(image, ks=5, iters=2):
    kernel = np.ones((ks, ks), np.uint8)
    dilate_image = cv2.dilate(image, kernel, iters)
    return dilate_image


def opening(image, ks=5, iters=2):
    kernel = np.ones((ks, ks), np.uint8)
    dilate_image = cv2.morphologyEx(image, cv2.MORPH_OPEN, kernel, iterations=iters)
    return dilate_image


def get_largest_contour(contours:list):
    if len(contours) == 1:
        ct = contours[0]
    else:
        candidate_regionprops = [cv2.contourArea(cn) for cn in contours]
        top_index = list(map(candidate_regionprops.index, heapq.nlargest(1, candidate_regionprops)))
        contours = [contours[i] for i in top_index]
        ct = np.concatenate(contours, axis=0)
    return ct


def get_largest_conneted_area(contours:list):
    ct = get_largest_contour(contours)
    hull = cv2.convexHull(ct)
    erode_max_x = hull[:, :, 0].max()
    erode_min_x = hull[:, :, 0].min()
    erode_max_y = hull[:, :, 1].max()
    erode_min_y = hull[:, :, 1].min()
    return (erode_min_x, erode_min_y, erode_max_x, erode_max_y)



def get_largest_cnt_from_mask(mask:np.ndarray):
    contours, hierarchy = cv2.findContours(mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)  # 轮廓提取
    cnt = get_largest_contour(contours)
    return cnt


def remove_small_objects(mask:np.ndarray):
    cnt = get_largest_cnt_from_mask(mask)
    return draw_contours_to_mask(mask, [cnt])


def segment_largest_roi_from_image(image:np.ndarray, debug=False):
    mask:np.ndarray = np.uint8(image > 0) * 255

    #retval, mask = cv2.threshold(image, 0, 255, cv2.THRESH_OTSU)

    h, w = mask.shape
    # if debug:
    #     cv_show('mask', mask, resize=0.7)
    #     cv2.waitKey(0)


    #erode_mask = erode(mask, ks=100, iters=10)
    dilate_mask = opening(mask, ks=20, iters=2)
    if debug:
        cv_show('dilate_mask', np.concatenate((image, mask, dilate_mask), axis=1), resize=0.7)
        cv2.waitKey(0)

    contours, hierarchy = cv2.findContours(dilate_mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)  # 轮廓提取
    dilate_min_x, dilate_min_y, dilate_max_x, dilate_max_y = get_largest_conneted_area(contours)
    roi_xyxy =  np.array([dilate_min_x, dilate_min_y, dilate_max_x, dilate_max_y])

    roi_xyxy[0::2] = np.clip(roi_xyxy[0::2], 0, w)
    roi_xyxy[1::2] = np.clip(roi_xyxy[1::2], 0, h)

    roi_xywh = np.array([roi_xyxy[0], roi_xyxy[1], roi_xyxy[2]-roi_xyxy[0], roi_xyxy[3]-roi_xyxy[1]])

    roi_w_ratio = roi_xywh[2] / w
    roi_h_ratio = roi_xywh[3] / h

    defalut_offset_w = 10
    defalut_offset_h = 10
    if roi_h_ratio < 0.8:
        defalut_offset_h = int((h - roi_xywh[3]) / 4) +50

    roi_xyxy[0] = roi_xyxy[0] - defalut_offset_w
    roi_xyxy[1] = roi_xyxy[1] - defalut_offset_h
    roi_xyxy[2] = roi_xyxy[2] + defalut_offset_w
    roi_xyxy[3] = roi_xyxy[3] + defalut_offset_h

    roi_xyxy[0::2] = np.clip(roi_xyxy[0::2], 0, w)
    roi_xyxy[1::2] = np.clip(roi_xyxy[1::2], 0, h)

    return roi_xyxy.tolist()



def resize_uniform_same_size(image:np.ndarray, roi_rect:list, resize:float):
    #1. crop roi
    #2. resize roi
    #3.put roi into src image
    src_h = image.shape[0]
    src_w = image.shape[1]

    roi = image[roi_rect[1]:roi_rect[1]+roi_rect[3], roi_rect[0]:roi_rect[0]+roi_rect[2]]

    roi = cv2.resize(roi, (int(roi_rect[2]*resize), int(roi_rect[3]*resize)))

    src_center = (roi_rect[0] + roi_rect[2]//2, roi_rect[1] + roi_rect[3]//2)


    h, w = roi.shape

    resize_roi_rect = (src_center[0] - w//2, src_center[1] - h//2, w, h)

    resize_image = np.zeros_like(image)

    resize_image[resize_roi_rect[1]:resize_roi_rect[1]+resize_roi_rect[3], resize_roi_rect[0]:resize_roi_rect[0]+resize_roi_rect[2]] =  roi

    return resize_image


def find_bounding_rect(mask):
    contours, hierarchy = cv2.findContours(mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)  # 轮廓提取
    ct = get_largest_contour(contours)
    x, y, w, h = cv2.boundingRect(ct)
    roi_xyxy = np.array([x, y, x+w, y+h])
    h, w = mask.shape[0], mask.shape[1]
    roi_xyxy[0::2] = np.clip(roi_xyxy[0::2], 0, w)
    roi_xyxy[1::2] = np.clip(roi_xyxy[1::2], 0, h)
    return roi_xyxy, ct
    #return np.array([roi_xyxy[0], roi_xyxy[1], roi_xyxy[2] - roi_xyxy[0], roi_xyxy[3] - roi_xyxy[1]], dtype=np.int)


def main():
    image_path = '/home/blake/data/medical/datasets/vertebral/2020.03.08/baimayixi+m+13/bmyx1.jpg'
    # image = load_image_obj(str(image_path), 'pydcm')

    data_path = Path("/home/blake/data/medical/datasets/vertebral/2020.03.08")
    image_items = sorted(data_path.rglob('*.jpg'))
    for image_path in image_items:
        image = cv2.imread(str(image_path), cv2.IMREAD_GRAYSCALE)
        min_x, min_y, max_x, max_y = segment_largest_from_image(image, debug=False)
        show_mask = np.repeat(image[..., np.newaxis], 3, 2)
        show_mask = cv2.rectangle(show_mask, (min_x, min_y), (max_x, max_y), (0, 255, 0), 5)

        cv_show('show', show_mask[min_y:max_y, min_x:max_x], 0.5)
        cv2.waitKey(0)

if __name__ == '__main__':
    main()

